package com.learn.grpc.loadbalance;

import io.grpc.LoadBalancer;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class LbSubChannelPicker extends LoadBalancer.SubchannelPicker {

    // 轮询索引
    private final AtomicInteger roundIdx = new AtomicInteger();

    // 子 channel 列表
    private List<LoadBalancer.Subchannel> subChannelList;

    public LbSubChannelPicker(List<LoadBalancer.Subchannel> subChannelList) {
        this.subChannelList = subChannelList;
    }

    @Override
    public LoadBalancer.PickResult pickSubchannel(LoadBalancer.PickSubchannelArgs args) {
        // subChannel 为空，则返回空的 Result
        if (subChannelList == null) {
            return LoadBalancer.PickResult.withNoResult();
        }

        // 轮询结束后，重置索引
        if (roundIdx.get() >= subChannelList.size()) {
            roundIdx.set(0);
        }

        // 轮询获取下一个 channel
        LoadBalancer.Subchannel selectNextChannel = subChannelList.get(roundIdx.getAndIncrement());

        return LoadBalancer.PickResult.withSubchannel(selectNextChannel);
    }


}
